import os
import requests
from dotenv import load_dotenv
from functools import wraps
from typing import Annotated
from sec_api import ExtractorApi, QueryApi, RenderApi
from fastmcp import FastMCP
import asyncio
from fastmcp import Client

# 加载环境变量
load_dotenv()

# 初始化FastMCP
mcp = FastMCP('sec_utils.mcp')

# 常量定义
CACHE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cache")
PDF_GENERATOR_API = "https://api.sec-api.io/filing-reader"
FMP_BASE = "https://financialmodelingprep.com"
SEC_BASE = "https://data.sec.gov"

# 全局API客户端
extractor_api = None
query_api = None
render_api = None

# 初始化SEC API
def init_sec_api(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        global extractor_api, query_api, render_api
        if os.environ.get("SEC_API_KEY") is None:
            return "请设置环境变量SEC_API_KEY以使用sec_api"
        
        if not extractor_api or not query_api or not render_api:
            extractor_api = ExtractorApi(os.environ["SEC_API_KEY"])
            query_api = QueryApi(os.environ["SEC_API_KEY"])
            render_api = RenderApi(os.environ["SEC_API_KEY"])
        
        return func(*args, **kwargs)
    return wrapper

SavePathType = Annotated[str, "File path to save data. If None, data is not saved."]
def decorate_all_methods(decorator):
    def class_decorator(cls):
        for attr_name, attr_value in cls.__dict__.items():
            if callable(attr_value):
                setattr(cls, attr_name, decorator(attr_value))
        return cls

    return class_decorator
# 获取SEC报告
@init_sec_api
def get_sec_report(
    ticker_symbol: Annotated[str, "股票代码，例如AAPL、MSFT"],
    fyear: Annotated[str, "财政年度，默认为'latest'表示最新"] = "latest",
) -> str:
    """
    获取给定股票和年份的10-K报告的URL和提交日期
    
    :param ticker_symbol: 股票代码
    :param fyear: 财政年度，默认为'latest'获取最新报告
    :result: 包含报告链接和提交日期的字符串
    """
    # FMP API Key
    try:
        api_key = fmp_api_key
    except NameError:
        api_key = os.getenv("FMP_API_KEY", "")

    # SEC 要求的User-Agent
    sec_ua = os.getenv("SEC_UA", "FinRobot/0.1 (you@example.com)")
    SEC_HEADERS = {"User-Agent": sec_ua}

    # 获取CIK
    def _get_cik(symbol: str) -> str | None:
        # 1) 尝试FMP API
        try:
            url = f"{FMP_BASE}/stable/sec-profile"
            params = {"symbol": symbol, "apikey": api_key}
            r = requests.get(url, params=params, timeout=30)
            if r.status_code == 200:
                data = r.json()
                rec = data[0] if isinstance(data, list) and data else data
                cik = rec.get("cik")
                if cik:
                    return str(cik).zfill(10)
        except Exception:
            pass
        
        # 2) 尝试SEC的ticker→CIK映射
        try:
            m = requests.get(f"{SEC_BASE}/files/company_tickers.json",
                            headers=SEC_HEADERS, timeout=30).json()
            for _, entry in m.items():
                if entry.get("ticker", "").upper() == symbol.upper():
                    return str(entry.get("cik_str")).zfill(10)
        except Exception:
            pass
        return None

    # 组装SEC文档链接
    def _build_sec_urls(cik_10: str, accession: str, primary_doc: str) -> tuple[str, str]:
        cik_nozero = str(int(cik_10))          # 去除前导0
        acc_nodash = accession.replace("-", "")
        doc_url = f"https://www.sec.gov/Archives/edgar/data/{cik_nozero}/{acc_nodash}/{primary_doc}"
        index_url = f"https://www.sec.gov/Archives/edgar/data/{cik_nozero}/{acc_nodash}/{accession}-index.htm"
        return doc_url, index_url

    # 主流程
    cik = _get_cik(ticker_symbol)
    if not cik:
        return f"无法获取{ticker_symbol}的CIK编码"

    subm_url = f"{SEC_BASE}/submissions/CIK{cik}.json"
    r = requests.get(subm_url, headers=SEC_HEADERS, timeout=30)
    if r.status_code != 200:
        return f"获取SEC数据失败: {r.status_code} - {r.text[:200]}"

    try:
        j = r.json()
        recent = j.get("filings", {}).get("recent", {})
        forms = recent.get("form", [])
        filing_dates = recent.get("filingDate", [])
        accessions = recent.get("accessionNumber", [])
        primary_docs = recent.get("primaryDocument", [])
    except Exception as e:
        return f"解析SEC响应失败: {e}"

    # 选择10-K报告
    idx = None
    if fyear == "latest":
        for i, f in enumerate(forms):
            if f == "10-K":
                idx = i
                break
    else:
        target = str(fyear)
        for i, (f, d) in enumerate(zip(forms, filing_dates)):
            if f == "10-K" and isinstance(d, str) and d[:4] == target:
                idx = i
                break

    if idx is None:
        return f"{ticker_symbol}在{fyear}年度没有找到10-K报告"

    accession = accessions[idx]
    filing_date = filing_dates[idx]
    primary_doc = primary_docs[idx]
    doc_url, index_url = _build_sec_urls(cik, accession, primary_doc)

    return f"Link: {doc_url or index_url}\nFiling Date: {filing_date}"

# 注册为FastMCP工具
@mcp.tool()
def get_10k_metadata(
    ticker: Annotated[str, "股票代码，例如AAPL、MSFT"],
    start_date: Annotated[str, "搜索范围的开始日期，格式为yyyy-mm-dd"],
    end_date: Annotated[str, "搜索范围的结束日期，格式为yyyy-mm-dd"],
) -> dict:
    """
    在给定时间段内搜索10-K文件，并返回最新的元数据
    
    :param ticker: 股票代码
    :param start_date: 开始日期，格式为yyyy-mm-dd
    :param end_date: 结束日期，格式为yyyy-mm-dd
    :result: 10-K报告的元数据字典，如果未找到则返回None
    """
    @init_sec_api
    def inner():
        query = {
            "query": f'ticker:"{ticker}" AND formType:"10-K" AND filedAt:[{start_date} TO {end_date}]',
            "from": 0,
            "size": 10,
            "sort": [{"filedAt": {"order": "desc"}}],
        }
        response = query_api.get_filings(query)
        if response["filings"]:
            return response["filings"][0]
        return None
    
    return inner()

#该函数和get_10k_metadata函数相同，为了方便后续函数调用
def get_10k_metadata1(
    ticker: Annotated[str, "股票代码，例如AAPL、MSFT"],
    start_date: Annotated[str, "搜索范围的开始日期，格式为yyyy-mm-dd"],
    end_date: Annotated[str, "搜索范围的结束日期，格式为yyyy-mm-dd"],
) -> dict:
    """
    在给定时间段内搜索10-K文件，并返回最新的元数据
    
    :param ticker: 股票代码
    :param start_date: 开始日期，格式为yyyy-mm-dd
    :param end_date: 结束日期，格式为yyyy-mm-dd
    :result: 10-K报告的元数据字典，如果未找到则返回None
    """
    @init_sec_api
    def inner():
        query = {
            "query": f'ticker:"{ticker}" AND formType:"10-K" AND filedAt:[{start_date} TO {end_date}]',
            "from": 0,
            "size": 10,
            "sort": [{"filedAt": {"order": "desc"}}],
        }
        response = query_api.get_filings(query)
        if response["filings"]:
            return response["filings"][0]
        return None
    
    return inner()

@mcp.tool()
def download_10k_filing(
    ticker: Annotated[str, "股票代码"],
    start_date: Annotated[
        str, "10-K 文件搜索范围的起始日期，格式为 yyyy-mm-dd"
    ],
    end_date: Annotated[
        str, "10-K 文件搜索范围的结束日期，格式为 yyyy-mm-dd"
    ],
    save_folder: Annotated[
        str, "用于保存下载文件的文件夹名称"
    ],
) -> str:
    """在指定时间范围内下载给定股票的最新 10-K 文件（htm 格式）。"""
    metadata = get_10k_metadata1(ticker, start_date, end_date)
    if metadata:
        ticker = metadata["ticker"]
        url = metadata["linkToFilingDetails"]

        try:
            date = metadata["filedAt"][:10]
            file_name = date + "_" + metadata["formType"] + "_" + url.split("/")[-1]

            if not os.path.isdir(save_folder):
                os.makedirs(save_folder)

            file_content = render_api.get_filing(url)
            file_path = os.path.join(save_folder, file_name)
            with open(file_path, "w") as f:
                f.write(file_content)
            return f"{ticker}: 下载成功。已保存至 {file_path}"
        except:
            return f"❌ {ticker}: 下载失败：{url}"
    else:
        return f"未找到 {ticker} 在 2023 年的 10-K 文件"



@mcp.tool()
def download_10k_pdf(
    ticker: Annotated[str, "股票代码"],
    start_date: Annotated[
        str, "10-K 文件搜索范围的起始日期，格式为 yyyy-mm-dd"
    ],
    end_date: Annotated[
        str, "10-K 文件搜索范围的结束日期，格式为 yyyy-mm-dd"
    ],
    save_folder: Annotated[
        str, "用于保存下载的 PDF 文件的文件夹名称"
    ],
) -> str:
    """在指定时间范围内下载给定股票的最新 10-K 文件（PDF 格式）。"""
    metadata = get_10k_metadata1(ticker, start_date, end_date)
    if metadata:
        ticker = metadata["ticker"]
        filing_url = metadata["linkToFilingDetails"]

        try:
            date = metadata["filedAt"][:10]
            print(filing_url.split("/")[-1])
            file_name = (
                date
                + "_"
                + metadata["formType"].replace("/A", "")
                + "_"
                + filing_url.split("/")[-1]
                + ".pdf"
            )

            # 若保存文件夹不存在，则创建
            if not os.path.isdir(save_folder):
                os.makedirs(save_folder)

            # 构建 PDF 下载 API 地址
            api_url = f"{PDF_GENERATOR_API}?token={os.environ['SEC_API_KEY']}&type=pdf&url={filing_url}"
            response = requests.get(api_url, stream=True)
            response.raise_for_status()

            # 以二进制方式写入 PDF 文件
            file_path = os.path.join(save_folder, file_name)
            with open(file_path, "wb") as file:
                for chunk in response.iter_content(chunk_size=8192):
                    file.write(chunk)
            return f"{ticker}: 下载成功。已保存至 {file_path}"
        except Exception as e:
            return f"❌ {ticker}: 下载失败：{filing_url}, 错误信息：{e}"
    else:
        return f"未找到 {ticker} 在 2023 年的 10-K 文件"


@mcp.tool()
def get_10k_section(
    ticker_symbol: Annotated[str, "股票代码"],
    fyear: Annotated[str, "10-K 报告对应的会计年度"],
    section: Annotated[
        str | int,
        "要提取的 10-K 报告章节，应在 [1, 1A, 1B, 2, 3, 4, 5, 6, 7, 7A, 8, 9, 9A, 9B, 10, 11, 12, 13, 14, 15] 范围内",
    ],
    report_address: Annotated[
        str,
        "10-K 报告的网页链接；如果未指定，则会通过 fmp API 自动获取报告地址",
    ] = None,
    save_path: SavePathType = None,
) -> str:
    """
    从 SEC API 获取指定公司的 10-K 报告的特定章节内容。
    """
    if isinstance(section, int):
        section = str(section)
    if section not in [
        "1A",
        "1B",
        "7A",
        "9A",
        "9B",
    ] + [str(i) for i in range(1, 16)]:
        raise ValueError(
            "章节参数必须在 [1, 1A, 1B, 2, 3, 4, 5, 6, 7, 7A, 8, 9, 9A, 9B, 10, 11, 12, 13, 14, 15] 范围内"
        )

    # 可选缓存路径示例
    # os.makedirs(f"{self.project_dir}/10k", exist_ok=True)
    # report_name = f"{self.project_dir}/10k/section_{section}.txt"
    # 如果启用缓存且文件存在，则直接读取缓存内容
    # if USE_CACHE and os.path.exists(report_name):
    #     with open(report_name, "r") as f:
    #         section_text = f.read()
    # else:

    # 若未提供报告地址，则通过 API 获取
    if report_address is None:
        report_address = get_sec_report(ticker_symbol, fyear)
        if report_address.startswith("Link: "):
            report_address = report_address.lstrip("Link: ").split()[0]
        else:
            return report_address  # 调试信息

    # 定义缓存路径
    cache_path = os.path.join(
        CACHE_PATH, f"sec_utils/{ticker_symbol}_{fyear}_{section}.txt"
    )

    # 若缓存文件存在，则直接读取
    if os.path.exists(cache_path):
        with open(cache_path, "r") as f:
            section_text = f.read()
    else:
        print(">>> 调试信息：正在从 SEC 获取章节内容", report_address, section)
        section_text = extractor_api.get_section(report_address, section, "text")
        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
        with open(cache_path, "w") as f:
            f.write(section_text)

    # 如果指定了保存路径，则保存提取结果
    if save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        with open(save_path, "w") as f:
            f.write(section_text)

    return section_text


GET_10K_METADATA_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_10k_metadata",
        "description": "在给定时间范围内搜索 10-K，并返回最新一条的元数据。",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "股票代码，例如 AAPL、MSFT。"
                },
                "start_date": {
                    "type": "string",
                    "description": "搜索范围开始日期（YYYY-MM-DD）。",
                    "format": "date"
                },
                "end_date": {
                    "type": "string",
                    "description": "搜索范围结束日期（YYYY-MM-DD）。",
                    "format": "date"
                }
            },
            "required": ["ticker", "start_date", "end_date"]
        }
    }
}

DOWNLOAD_10K_FILING_SCHEMA = {
    "type": "function",
    "function": {
        "name": "download_10k_filing",
        "description": "下载指定股票在时间范围内最新的 10-K（HTML/HTM 原文）到本地文件夹。",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "股票代码，例如 AAPL、MSFT。"
                },
                "start_date": {
                    "type": "string",
                    "description": "10-K 搜索起始日期（YYYY-MM-DD）。",
                    "format": "date"
                },
                "end_date": {
                    "type": "string",
                    "description": "10-K 搜索结束日期（YYYY-MM-DD）。",
                    "format": "date"
                },
                "save_folder": {
                    "type": "string",
                    "description": "保存下载文件的文件夹路径（不存在将自动创建）。"
                }
            },
            "required": ["ticker", "start_date", "end_date", "save_folder"]
        }
    }
}

DOWNLOAD_10K_PDF_SCHEMA = {
    "type": "function",
    "function": {
        "name": "download_10k_pdf",
        "description": "下载指定股票在时间范围内最新的 10-K（PDF 转换版）到本地文件夹。",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "股票代码，例如 AAPL、MSFT。"
                },
                "start_date": {
                    "type": "string",
                    "description": "10-K 搜索起始日期（YYYY-MM-DD）。",
                    "format": "date"
                },
                "end_date": {
                    "type": "string",
                    "description": "10-K 搜索结束日期（YYYY-MM-DD）。",
                    "format": "date"
                },
                "save_folder": {
                    "type": "string",
                    "description": "保存下载 PDF 的文件夹路径（不存在将自动创建）。"
                }
            },
            "required": ["ticker", "start_date", "end_date", "save_folder"]
        }
    }
}

GET_10K_SECTION_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_10k_section",
        "description": "提取指定公司 10-K 报告的特定章节文本（如 1A、7、7A、9A 等）。",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker_symbol": {
                    "type": "string",
                    "description": "股票代码，例如 AAPL、MSFT。"
                },
                "fyear": {
                    "type": "string",
                    "description": "10-K 对应会计年度（四位年份，例如 2024）。",
                    "pattern": "^\\d{4}$"
                },
                "section": {
                    "type": ["string", "integer"],
                    "description": "要提取的章节。可为整数或字符串；有效值包括 [1, 1A, 1B, 2, 3, 4, 5, 6, 7, 7A, 8, 9, 9A, 9B, 10, 11, 12, 13, 14, 15]。",
                    "enum": ["1","1A","1B","2","3","4","5","6","7","7A","8","9","9A","9B","10","11","12","13","14","15"]
                },
                "report_address": {
                    "type": ["string", "null"],
                    "description": "10-K 报告网页链接（可选；未提供时会自动获取最新链接）。",
                    "format": "uri"
                },
                "save_path": {
                    "type": ["string", "null"],
                    "description": "将提取结果另存为文本的完整文件路径（可选）。"
                }
            },
            "required": ["ticker_symbol", "fyear", "section"]
        }
    }
}

async def main():
    client = Client(mcp)
    async with client:
        # 查看可用工具
        tools = await client.list_tools()
        print('可用工具:', tools)
        
        # 示例：获取MSFT的10-K元数据
        ticker = "MSFT"
        start_date = "2024-01-01"
        end_date = "2025-01-01"
        
        print(f"\n获取{ticker}的10-K元数据...")
        # 调用工具并直接使用结果（不通过.result属性）
        result = await client.call_tool('get_10k_metadata', {
            'ticker': ticker,
            'start_date': start_date,
            'end_date': end_date
        })
        print("10-K元数据:", result)
        
        # 如果找到元数据，提取1A章节
        if result.data:
            filing_details = result.data
            print(f"\n提取{ticker}的10-K报告1A章节...")
            section_text = await client.call_tool('get_10k_section', {
                'ticker_symbol': ticker,
                'fyear': "2024",
                'section': "1A",
                'report_address': filing_details["linkToFilingDetails"]
            })
            print("10-K Item 1A 预览:", section_text)
        
        # 示例：下载MSFT的10-K报告
        print(f"\n下载MSFT的10-K报告...")
        download_result = await client.call_tool('download_10k_filing', {
            'ticker': 'MSFT',
            'start_date': start_date,
            'end_date': end_date,
            'save_folder': './downloaded_reports'
        })


if __name__ == '__main__':
    asyncio.run(main())
